import torch
from .conv import Conv, DWConv
from .block import RepVGGDW



class DCB(torch.nn.Module):
    def __init__(self, c1: int, c2: int):
        super().__init__()
        self.cv1 = Conv(c1, c2, 3, 1)
        self.dwconv = DWConv(c2, c2, 1, 1)
        self.cv2 = Conv(c2, c2, 1, 1)
        self.repvggdw = RepVGGDW(c2)
        self.cv3 = Conv(c2, c2, 1, 1)

    def forward(self, x):
        branch1 = self.cv1(x)
        branch2 = self.cv3(self.repvggdw(self.cv2(self.dwconv(branch1))))
        return branch1 + branch2


class C2fDCB(torch.nn.Module):
    def __init__(self, c1: int, c2: int, n: int = 1):      # channel_in, channel_out, repeat_num
        super().__init__()
        self.cv1 = Conv(c1, c2, 3, 1)
        # 重命名 modules 为 blocks 以避免命名冲突
        self.blocks = torch.nn.Sequential(*[DCB(c2, c2) for _ in range(n)])
        # 调整通道数，将n个c2通道合并为c2通道
        self.cv2 = Conv(c2 * n, c2, 1, 1) if n > 1 else torch.nn.Identity()

    def forward(self, x):
        x = self.cv1(x)  # 初始卷积
        outputs = []
        current = x

        # 迭代重命名后的 blocks 属性
        for block in self.blocks:
            current = block(current)
            outputs.append(current)

        # 拼接所有输出并调整通道数
        return self.cv2(torch.cat(outputs, dim=1))

